# 导入tushare
import os
import platform
import threading
from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED
from datetime import date

import tushare as ts
import pandas as pd
import os


class StockDataPrepare:
    # 初始化pro接口
    pro = ts.pro_api('591d1fee4ba2ebe860c911a5e58bd4d344d9458c6a9e16b09b5afd46')
    basepath=os.path.split(os.path.realpath(__file__))[0]
    hdf5path = basepath+'/HDF5DATA/'
    # getStock_list from tushare
    def getStockListFromTushar(self):
        # 拉取数据
        df = self.pro.stock_basic(**{
            "ts_code": "",
            "name": "",
            "exchange": "",
            "market": "",
            "is_hs": "",
            "list_status": "",
            "limit": "",
            "offset": ""
        }, fields=[
            "ts_code",
            "symbol",
            "name",
            "area",
            "industry",
            "market",
            "list_date",
            "fullname",
            "exchange"
        ])
        # print(type(df))
        # print(df.shape)
        df.to_csv('stock_list.csv')
    # get stock list for file
    def getStock_list_from_csvFile(self):
        df = pd.read_csv(self.hdf5path +'../stock_list.csv')
        return df
        # for index,row in df.iterrows():
        #     # print(index,row)
        #     print(index, row.ts_code)
    def get_limit_from_tushare(self, code_markect, limit='', start_date="", end_date=''):
        df = self.pro.stk_limit(**{
            "ts_code": code_markect,
            "trade_date": "",
            "start_date": "",
            "end_date": "",
            "offset": "",
            "limit": ""
        }, fields=[
            "trade_date",
            "ts_code",
            "up_limit",
            "down_limit",
            "pre_close"
        ])
        print(df)
    # 每分钟内最多调取500次，每次5000条数据，相当于单次提取23年历史
    # 获取某一股票所有数据
    def getKline_from_tushare(self,code_markect,limit='',start_date="",end_date=''):
        # print(code_markect,limit,start_date,end_date)
        df = self.pro.daily(**{
            "ts_code": code_markect,
            "trade_date": "",
            "start_date": start_date,
            "end_date": end_date,
            "offset": "",
            "limit": limit
        }, fields=[
            "ts_code",
            "trade_date",
            "open",
            "high",
            "low",
            "close",
            "pre_close",
            "change",
            "pct_chg",
            "vol",
            "amount"
        ])
        # print(df)
        # df0 = df.rename(columns={'trade_date': 'timetag', 'vol': 'volumn'})
        # print(df0)
        # df0 = df0[['timetag', 'open', 'high', 'low', 'close', 'volumn', 'amount']]
        # df0 = df0.iloc[::-1]
        return df


    def get_klines_from_hdf5(self,stock_code):
        # GUI command: vitables
        [code,markect] = stock_code.split('.')
        # print(code,markect)
        # num= int(code)%10

        store = pd.HDFStore(self.hdf5path + markect +code + '.h5')
        # print(store.keys())
        df =store['df']
        store.close()
        # print(df)
        # return df['trade_date'][0]
        return df
    def set_klines_to_hdf5(self,stock_code,df):
        [code,markect] = stock_code.split('.')
        # print('./HDF5Date/' + markect + code + '.h5')
        try:
            # store = pd.HDFStore(self.hdf5path + markect + code + '.h5',format='table')
            df.to_hdf(self.hdf5path + markect + code + '.h5', key='df', mode='w', format='table')

        except Exception as e:
            print(e)

    # todo 一天天去更新
    def update_allstock_onday(self, startDay):
        # current_date = date.today()
        # today = current_date.strftime('%Y%m%d')
        max_workers = 40
        pool = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix='测试线程')
        df0 = self.getKline_from_tushare('', '', startDay,startDay)
        print('df0.shape:',df0.shape)
        tasklist = []
        for index, row in df0.iterrows():
            [code, markect] = row.ts_code.split('.')
            if markect == 'BJ':
                continue
            print(index,row.ts_code)
            code_markect=row.ts_code
            try:
                task=pool.submit(self.update_one_stock,code_markect,row)
            except Exception as e:
                print(e)
            if (len(tasklist) > max_workers):
                wait(tasklist, return_when=ALL_COMPLETED)
                tasklist.clear()
            tasklist.append(task)
            # break
        wait(tasklist, return_when=ALL_COMPLETED)
        print('------------主线程执行结束----')
    def fix_data(self):
        df1 = self.get_klines_from_hdf5('689009.SH')
        df1.drop(df1.index[0], inplace=True)  # 删除第1行
        # print(df1)
        self.set_klines_to_hdf5('689009.SH', df1)
        return

        df= self.getStock_list_from_csvFile()
        for index, row in df.iterrows():
            if row.ts_code=='000001.SZ':
                continue
            print(row.ts_code)
            [code, markect] = row.ts_code.split('.')

            if markect == 'BJ':
                continue
            df1= self.get_klines_from_hdf5(row.ts_code)
            df1.drop(df1.index[0], inplace=True)    # 删除第1行
            # print(df1)
            self.set_klines_to_hdf5(row.ts_code,df1)
            # df0 = self.get_klines_from_hdf5(row.ts_code)
            # print(df0)
            # break
    def update_one_stock(self,code_markect,row):
        # print('in update_one_stock')
        try:
            df = self.get_klines_from_hdf5(code_markect)
        except Exception as e:
            print(e)
        # print(df)
        # row.drop('ts_code', inplace=True)
        # df = df.append(row, ignore_index=True)
        # print(code_markect)
        if len(df) == 0:
            df = self.getKline_from_tushare(code_markect)
        else:
            # print('guava')
            # print(df.iloc[-1].trade_date , row.trade_date)
            if df.iloc[0].trade_date == row.trade_date:
                # print(df.iloc[-1].tolist())
                print(code_markect, 'continue')
                return
        df = pd.concat([pd.DataFrame([row]),df.loc[::]], axis=0, ignore_index=True)
        # print(df)
        # print(index, row.ts_code)
        # self.set_klines_to_hdf5(code_markect, df)
        [code, markect] = code_markect.split('.')
        try:
            df.to_hdf(self.hdf5path + markect + code + '.h5', key='df', mode='w', format='table')
        except Exception as e:
            print(e)
    def autoUpdata(self):
        sl = StockDataPrepare()
        dfo= sl.get_klines_from_hdf5('000004.SZ')
        print(dfo)
        df = sl.get_days_trade('20220510', '20220515')
        # get last day
        # get today
        return
        for index, item in df.iterrows():
            print(item.cal_date, item.is_open)
            if item.is_open == 1:
                sl.update_allstock_onday(item.cal_date)
            else:
                print(type(item.is_open))

    def first_time_getKline_from_tushare(self):
        df = self.getStock_list_from_csvFile()
        for index, row in df.iterrows():
            [code, markect] = row.ts_code.split('.')
            if markect=='BJ':
                continue
            # print(index,row)
            print(index, row.ts_code)
            df=self.getKline_from_tushare(row.ts_code,start_date='')
            self.set_klines_to_hdf5(row.ts_code,df)
    def getklineFromHDF5_toSingleFile(self):
        threads = []
        for i in range(0,10):
            try:
                t=threading.Thread(target=threadOneENDFile,args=('SH_'+str(i),))
                print('create')

                threads.append(t)
                t.start()
            except Exception as e:
                print('can not star thread')
                print(e)
        # 等待所有子线程结束再运行主线程
        [thread.join() for thread in threads]
    def get_files_to_singleH5(self):
        for f in os.listdir(self.hdf5path):
            print(f)
            # store = pd.HDFStore(self.hdf5path + f)
            # df = store['df']
            # df.to_hdf('./HDF5DATA/'+f, key='df', mode='w', format='table')
            # store.close()
    def get_days_trade(self,start,end):
        df = self.pro.trade_cal(**{
            "exchange": "",
            "cal_date": "",
            "start_date": start,
            "end_date": end,
            "is_open": "",
            "limit": "",
            "offset": ""
        }, fields=[
            "exchange",
            "cal_date",
            "is_open",
            "pretrade_date"
        ])
        return df
def threadOneENDFile(filename):
    hdf5path = 'H:/HDF5Date/'
    try:
        store = pd.HDFStore(hdf5path + filename + '.h5',complevel=9,complib="blosc:lz4")
        print(store.keys())
    except Exception as e:
        print(e)
    for item in store.keys():
        # print(item.replace('/', ''))
        df = store[item]
        store1 = pd.HDFStore(hdf5path+'single/' + item + '.h5')
        store1['df'] = df
        store1.close()
    store.close()


